import sys
import os
from os import system,name
import time
if name == 'nt':
	from Cryptodome.PublicKey import RSA
	from Cryptodome.Random import get_random_bytes
	from Cryptodome.Cipher import AES, PKCS1_OAEP
else:
	from Crypto.PublicKey import RSA
	from Crypto.Random import get_random_bytes
	from Crypto.Cipher import AES, PKCS1_OAEP

#Just for clear terminal
def clear():
	system('cls') if name == 'nt' else system('clear')

def main():
	clear()
	print("1. Generate keys")
	print("2. Encrypt File")
	print("3. Decrypt File")
	print("4. Exit")
	main_input = input(">>> ")
	if main_input == "1":
		enc_key = RSA.generate(4096) #Generating RSA key with len = 4096. 
		privkeyname = input("Enter your private key name: ") #Entering name for your private key
		pubkeyname = input("Enter your public key name: ") #Entering name for your public key
		private_key = enc_key.export_key() #Exporting encrypted prime numbers from key
		prkey_out = open("Keys/"+ privkeyname + ".pem", "wb")
		prkey_out.write(private_key) #Writing private key to file
		prkey_out.close() #Closing your private key so you can instantly use him, without closing a program
		public_key = enc_key.publickey().export_key() #Exporting encrypted big number from private key based on first generated key
		pbkey_out = open("Keys/" + pubkeyname + ".pem", "wb")
		pbkey_out.write(public_key) #Writing public key to file
		pbkey_out.close()  #Closing your public key so you can instantly use him, without closing a program
		print("Done! You can find your keys in 'Keys' folder")
		print("Press [ENTER] to return to main menu")
		asdasddsad = input("")
		main()
	elif main_input == "2":
		file_ = input("Input your file to encrypt: ") #Entering your file_ to encrypt
		pubkey = input("Input your public key: ") #Entering your pubkey to encrypt file_ with him
		f = open(file_, "rb") #Opening your file_
		data = f.read(); f.close() #Reading data from your file_, closing it
		file_out = open(str(file_)+".bin", "wb") #Writing output file
		session_key = get_random_bytes(16) #Getting random bytes
		recipient_key = RSA.import_key(open(pubkey).read()) #Creating recipient key from pubkey
		cipher_rsa = PKCS1_OAEP.new(recipient_key) #creating cipher rsa from recipient key
		enc_session_key = cipher_rsa.encrypt(session_key) #Creating enc session key from random bytes
		cipher_aes = AES.new(session_key, AES.MODE_EAX) #Strengthen it all with AES encryption
		ciphertext, tag = cipher_aes.encrypt_and_digest(data) #Finally, encrypting the data from file_
		[ file_out.write(x) for x in (enc_session_key, cipher_aes.nonce, tag, ciphertext) ] #Writing all encrypted data to the new, encrypted file_
		f.close()
		file_out.close()
		print("File encrypted succesfully!")
		print("Press [ENTER] to return to main menu")
		asdasddsad = input("") #Just for hold the terminal, nevermind
		main()
	elif main_input == "3":
		privkey = input("Input your private key: ") #Entering your private key to decrypt 
		enc_file_ = input("Input your encrypted file: ") #Entering your encrypted file_ to decrypt
		file_in = open(enc_file_, "rb") 
		file_out = open(str(enc_file_[:-4]), "wb") #Opening your encrypted file_
		private_key = RSA.import_key(open(privkey).read()) #Reading your private key
		enc_session_key, nonce, tag, ciphertext = \
		   [ file_in.read(x) for x in (private_key.size_in_bytes(), 16, 16, -1) ] #Creating enc session key from private key
		cipher_rsa = PKCS1_OAEP.new(private_key)
		session_key = cipher_rsa.decrypt(enc_session_key) #Decrypting your data from encrypted file_
		cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
		data = cipher_aes.decrypt_and_verify(ciphertext, tag) #Verifying that is all right
		file_out.write(data) #Writing decrypted data to your file_
		file_in.close()
		file_out.close()
		print("File decrypted succesfully!")
		os.remove(enc_file_) #Removing encrypted file_ from your system, you don't need it anymore
		print("Press [ENTER] to return to main menu")
		asdasdasd = input("")
		main()
	elif main_input == "4":
		exit(0)
	else:
		print("Error input, repeating..")
		time.sleep(1)
		main()


if __name__ == "__main__":
	try:
		main()
	except KeyboardInterrupt:
		clear()
		print("Bye")
		exit(0)